
/**
  ******************************************************************************
  * Copyright 2021 The Microbee Authors. All Rights Reserved.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  * 
  * http://www.apache.org/licenses/LICENSE-2.0
  * 
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * 
  * @file       fms_attitude.c
  * @author     baiyang
  * @date       2023-7-13
  ******************************************************************************
  */

/*----------------------------------include-----------------------------------*/
#include "fms.h"

#include <uITC/uITC.h>
#include <uITC/uITC_msg.h>
/*-----------------------------------macro------------------------------------*/

/*----------------------------------typedef-----------------------------------*/

/*---------------------------------prototype----------------------------------*/

/*----------------------------------variable----------------------------------*/

/*-------------------------------------os-------------------------------------*/

/*----------------------------------function----------------------------------*/
/*
  calculate a new aerodynamic_load_factor and limit nav_roll_cd to
  ensure that the load factor does not take us below the sustainable
  airspeed
 */
void fms_update_load_factor(void)
{
    float demanded_roll = fabsf(fms.nav_roll_cd*0.01f);
    if (demanded_roll > 85) {
        // limit to 85 degrees to prevent numerical errors
        demanded_roll = 85;
    }
    fms.aerodynamic_load_factor = 1.0f / math_sqrtf(cosf(radians(demanded_roll)));

#if HAL_QUADPLANE_ENABLED
    if (quadplane.available() && quadplane.transition->set_FW_roll_limit(roll_limit_cd)) {
        nav_roll_cd = constrain_int32(nav_roll_cd, -roll_limit_cd, roll_limit_cd);
        return;
    }
#endif

    if (!fms.aparm.stall_prevention) {
        // stall prevention is disabled
        return;
    }
    if (fms_fly_inverted()) {
        // no roll limits when inverted
        return;
    }
#if HAL_QUADPLANE_ENABLED
    if (quadplane.tailsitter.active()) {
        // no limits while hovering
        return;
    }
#endif

    float max_load_factor = fms.smoothed_airspeed / MAX(fms.aparm.airspeed_min, 1);
    if (max_load_factor <= 1) {
        // our airspeed is below the minimum airspeed. Limit roll to
        // 25 degrees
        fms.nav_roll_cd = math_constrain_int32(fms.nav_roll_cd, -2500, 2500);
        fms.roll_limit_cd = MIN(fms.roll_limit_cd, 2500);
    } else if (max_load_factor < fms.aerodynamic_load_factor) {
        // the demanded nav_roll would take us past the aerodymamic
        // load limit. Limit our roll to a bank angle that will keep
        // the load within what the airframe can handle. We always
        // allow at least 25 degrees of roll however, to ensure the
        // aircraft can be maneuvered with a bad airspeed estimate. At
        // 25 degrees the load factor is 1.1 (10%)
        int32_t roll_limit = degrees(acosf(sq(1.0f / max_load_factor)))*100;
        if (roll_limit < 2500) {
            roll_limit = 2500;
        }
        fms.nav_roll_cd = math_constrain_int32(fms.nav_roll_cd, -roll_limit, roll_limit);
        fms.roll_limit_cd = MIN(fms.roll_limit_cd, roll_limit);
    }
}

/*
  adjust nav_pitch_cd for STAB_PITCH_DOWN_CD. This is used to make
  keeping up good airspeed in FBWA mode easier, as the plane will
  automatically pitch down a little when at low throttle. It makes
  FBWA landings without stalling much easier.
 */
void fms_adjust_nav_pitch_throttle(void)
{
    int8_t throttle = fms_throttle_percentage();
    if (throttle >= 0 && throttle < fms.aparm.throttle_cruise && fms.flight_stage != FWFS_VTOL) {
        float p = (fms.aparm.throttle_cruise - throttle) / (float)fms.aparm.throttle_cruise;
        fms.nav_pitch_cd -= fms.g.stab_pitch_down * 100.0f * p;
    }
}

/*
  return throttle percentage from 0 to 100 for normal use and -100 to 100 when using reverse thrust
 */
int8_t fms_throttle_percentage(void)
{
#if HAL_QUADPLANE_ENABLED
    if (quadplane.in_vtol_mode() && !quadplane.tailsitter.in_vtol_transition()) {
        return quadplane.motors->get_throttle_out() * 100.0;
    }
#endif
    float throttle = srv_channels_get_output_scaled(k_throttle);
    if (!fms_have_reverse_thrust()) {
        return math_constrain_int16(throttle, 0, 100);
    }
    return math_constrain_int16(throttle, -100, 100);
}

/*
  main stabilization function for all 3 axes
 */
void fms_stabilize()
{
    uint32_t now = time_millis();
#if HAL_QUADPLANE_ENABLED
    if (quadplane.available()) {
        quadplane.transition->set_FW_roll_pitch(nav_pitch_cd, nav_roll_cd);
    }
#endif

    if (now - fms.last_stabilize_ms > 2000) {
        // if we haven't run the rate controllers for 2 seconds then reset
        mode_reset_controllers(fms.control_mode);
    }
    fms.last_stabilize_ms = now;

    if (fms.control_mode == (mode_base_t)&fms.mode_training ||
        fms.control_mode == (mode_base_t)&fms.mode_manual) {
        mode_run(fms.control_mode);
    } else {
        mode_run(fms.control_mode);
    }

    uitc_sensor_baro sensor_baro = {0};
    itc_copy_from_hub(ITC_ID(sensor_baro), &sensor_baro);

    /*
      see if we should zero the attitude controller integrators. 
     */
    if (math_flt_zero(fms_get_throttle_input(false)) &&
        fabsf(fms.relative_altitude) < 5.0f && 
        fabsf(sensor_baro.velocity_ms) < 0.5f &&
        ahrs_get_speed_xy_cms(fms.ahrs) < 3) {
        // we are low, with no climb rate, and zero throttle, and very
        // low ground speed. Zero the attitude controller
        // integrators. This prevents integrator buildup pre-takeoff.
        rctrl_reset_i(&fms.roll_controller);
        pctrl_reset_i(&fms.pitch_controller);
        yctrl_reset_i(&fms.yaw_controller);

        // if moving very slowly also zero the steering integrator
        if (ahrs_get_speed_xy_cms(fms.ahrs) < 1) {
            sctrl_reset_i(&fms.steer_controller);
        }
    }
}

/*
  this gives the user control of the aircraft in stabilization modes
  using FBW style controls
 */
void fms_stabilize_stick_mixing_fbw()
{
    if (!fms_stick_mixing_enabled() ||
        fms.control_mode == (mode_base_t)&fms.mode_acro ||
        fms.control_mode == (mode_base_t)&fms.mode_fbwa ||
        fms.control_mode == (mode_base_t)&fms.mode_autotune ||
        fms.control_mode == (mode_base_t)&fms.mode_fbwb ||
        fms.control_mode == (mode_base_t)&fms.mode_cruise ||
#if HAL_QUADPLANE_ENABLED
        control_mode == &mode_qstabilize ||
        control_mode == &mode_qhover ||
        control_mode == &mode_qloiter ||
        control_mode == &mode_qland ||
        control_mode == &mode_qacro ||
#if QAUTOTUNE_ENABLED
        control_mode == &mode_qautotune ||
#endif
        !quadplane.allow_stick_mixing() ||
#endif  // HAL_QUADPLANE_ENABLED
        fms.control_mode == (mode_base_t)&fms.mode_training) {
        return;
    }
    // do FBW style stick mixing. We don't treat it linearly
    // however. For inputs up to half the maximum, we use linear
    // addition to the nav_roll and nav_pitch. Above that it goes
    // non-linear and ends up as 2x the maximum, to ensure that
    // the user can direct the plane in any direction with stick
    // mixing.
    float roll_input = rc_norm_input_dz(fms.channel_roll);
    if (roll_input > 0.5f) {
        roll_input = (3*roll_input - 1);
    } else if (roll_input < -0.5f) {
        roll_input = (3*roll_input + 1);
    }
    fms.nav_roll_cd += roll_input * fms.roll_limit_cd;
    fms.nav_roll_cd = math_constrain_int32(fms.nav_roll_cd, -fms.roll_limit_cd, fms.roll_limit_cd);

    if ((fms.control_mode == (mode_base_t)&fms.mode_loiter) && (fms_flight_option_enabled(FO_ENABLE_LOITER_ALT_CONTROL))) {
        // loiter is using altitude control based on the pitch stick, don't use it again here
        return;
    }

    float pitch_input = rc_norm_input_dz(fms.channel_pitch);
    if (pitch_input > 0.5f) {
        pitch_input = (3*pitch_input - 1);
    } else if (pitch_input < -0.5f) {
        pitch_input = (3*pitch_input + 1);
    }
    if (fms_fly_inverted()) {
        pitch_input = -pitch_input;
    }
    if (pitch_input > 0) {
        fms.nav_pitch_cd += pitch_input * fms.aparm.pitch_limit_max_cd;
    } else {
        fms.nav_pitch_cd += -(pitch_input * fms.pitch_limit_min_cd);
    }
    fms.nav_pitch_cd = math_constrain_int32(fms.nav_pitch_cd, fms.pitch_limit_min_cd, fms.aparm.pitch_limit_max_cd);
}

/*
  return true if the current settings and mode should allow for stick mixing
 */
bool fms_stick_mixing_enabled(void)
{
    if (!rcs_has_valid_input()) {
        // never stick mix without valid RC
        return false;
    }
#if AP_FENCE_ENABLED
    const bool stickmixing = fence_stickmixing();
#else
    const bool stickmixing = true;
#endif
#if HAL_QUADPLANE_ENABLED
    if (fms.control_mode == &fms.mode_qrtl &&
        quadplane.poscontrol.get_state() >= QuadPlane::QPOS_POSITION1) {
        // user may be repositioning
        return false;
    }
    if (quadplane.in_vtol_land_poscontrol()) {
        // user may be repositioning
        return false;
    }
#endif
    if (mode_does_auto_throttle(fms.control_mode) && mode_does_auto_navigation(fms.control_mode)) {
        // we're in an auto mode. Check the stick mixing flag
        if (fms.g.stick_mixing != SM_NONE &&
            fms.g.stick_mixing != SM_VTOL_YAW &&
            stickmixing) {
            return true;
        } else {
            return false;
        }
    }

    if (fms.failsafe.rc_failsafe && fms.g.fs_action_short == FS_ACTION_SHORT_FBWA) {
        // don't do stick mixing in FBWA glide mode
        return false;
    }

    // non-auto mode. Always do stick mixing
    return true;
}

/*
  this is the main roll stabilization function. It takes the
  previously set nav_roll calculates roll servo_out to try to
  stabilize the plane at the given roll
 */
void fms_stabilize_roll()
{
    if (fms_fly_inverted()) {
        // we want to fly upside down. We need to cope with wrap of
        // the roll_sensor interfering with wrap of nav_roll, which
        // would really confuse the PID code. The easiest way to
        // handle this is to ensure both go in the same direction from
        // zero
        fms.nav_roll_cd += 18000;
        if (fms.ahrs->roll_sensor_cd < 0) fms.nav_roll_cd -= 36000;
    }

    const float roll_out = fms_stabilize_roll_get_roll_out();
    srv_channels_set_output_scaled(k_aileron, roll_out);
}

float fms_stabilize_roll_get_roll_out()
{
    const float speed_scaler = fms.surface_speed_scaler;
#if HAL_QUADPLANE_ENABLED
    if (!quadplane.use_fw_attitude_controllers()) {
        // use the VTOL rate for control, to ensure consistency
        const auto &pid_info = quadplane.attitude_control->get_rate_roll_pid().get_pid_info();
        const float roll_out = rollController.get_rate_out(degrees(pid_info.target), speed_scaler);
        /* when slaving fixed wing control to VTOL control we need to decay the integrator to prevent
           opposing integrators balancing between the two controllers
        */
        rollController.decay_I();
        return roll_out;
    }
#endif

    bool disable_integrator = false;
    if (fms.control_mode == (mode_base_t)&fms.mode_stabilize && rc_get_control_in(fms.channel_roll) != 0) {
        disable_integrator = true;
    }
    return rctrl_get_servo_out(&fms.roll_controller, fms.nav_roll_cd - fms.ahrs->roll_sensor_cd, speed_scaler, disable_integrator,
                                        fms.ground_mode && !(fms_flight_option_enabled(FO_DISABLE_GROUND_PID_SUPPRESSION)));
}

/*
  this is the main pitch stabilization function. It takes the
  previously set nav_pitch and calculates servo_out values to try to
  stabilize the plane at the given attitude.
 */
void fms_stabilize_pitch()
{
    int8_t force_elevator = fms_takeoff_tail_hold();
    if (force_elevator != 0) {
        // we are holding the tail down during takeoff. Just convert
        // from a percentage to a -4500..4500 centidegree angle
        srv_channels_set_output_scaled(k_elevator, 45*force_elevator);
        return;
    }

    const float pitch_out = fms_stabilize_pitch_get_pitch_out();
    srv_channels_set_output_scaled(k_elevator, pitch_out);
}

float fms_stabilize_pitch_get_pitch_out()
{
    const float speed_scaler = fms.surface_speed_scaler;
#if HAL_QUADPLANE_ENABLED
    if (!quadplane.use_fw_attitude_controllers()) {
        // use the VTOL rate for control, to ensure consistency
        const auto &pid_info = quadplane.attitude_control->get_rate_pitch_pid().get_pid_info();
        const int32_t pitch_out = pitchController.get_rate_out(degrees(pid_info.target), speed_scaler);
        /* when slaving fixed wing control to VTOL control we need to decay the integrator to prevent
           opposing integrators balancing between the two controllers
        */
        pitchController.decay_I();
        return pitch_out;
    }
#endif
    // if LANDING_FLARE RCx_OPTION switch is set and in FW mode, manual throttle,throttle idle then set pitch to LAND_PITCH_CD if flight option FORCE_FLARE_ATTITUDE is set
#if HAL_QUADPLANE_ENABLED
    const bool quadplane_in_transition = quadplane.in_transition();
#else
    const bool quadplane_in_transition = false;
#endif

    int32_t demanded_pitch = fms.nav_pitch_cd + fms.aparm.pitch_trim_cd + srv_channels_get_output_scaled(k_throttle) * fms.g.kff_throttle_to_pitch;
    bool disable_integrator = false;
    if (fms.control_mode == (mode_base_t)&fms.mode_stabilize && rc_get_control_in(fms.channel_pitch) != 0) {
        disable_integrator = true;
    }
    /* force landing pitch if:
       - flare switch high
       - throttle stick at zero thrust
       - in fixed wing non auto-throttle mode
    */
#if 0
    if (!quadplane_in_transition &&
        !mode_is_vtol_mode(&fms.control_mode) &&
        !mode_does_auto_throttle(&fms.control_mode) &&
        fms.flare_mode == FM_ENABLED_PITCH_TARGET &&
        fms_throttle_at_zero()) {
        demanded_pitch = landing.get_pitch_cd();
    }
#endif

    return pctrl_get_servo_out(&fms.pitch_controller, demanded_pitch - fms.ahrs->pitch_sensor_cd, speed_scaler, disable_integrator,
                                         fms.ground_mode && !(fms_flight_option_enabled(FO_DISABLE_GROUND_PID_SUPPRESSION)));
}

/*
  stabilize the yaw axis. There are 3 modes of operation:

    - hold a specific heading with ground steering
    - rate controlled with ground steering
    - yaw control for coordinated flight    
 */
void fms_stabilize_yaw()
{
    bool ground_steering = false;

    if (landing_is_flaring()) {
        // in flaring then enable ground steering
        ground_steering = true;
    } else {
        // otherwise use ground steering when no input control and we
        // are below the GROUND_STEER_ALT
        ground_steering = (rc_get_control_in(fms.channel_roll) == 0 && 
                                            fabsf(fms.relative_altitude) < fms.g.ground_steer_alt);
        if (!landing_is_ground_steering_allowed()) {
            // don't use ground steering on landing approach
            ground_steering = false;
        }
    }


    /*
      first calculate steering for a nose or tail
      wheel. We use "course hold" mode for the rudder when either performing
      a flare (when the wings are held level) or when in course hold in
      FBWA mode (when we are below GROUND_STEER_ALT)
     */
    float steering_output = 0.0f;
    if (landing_is_flaring() ||
        (fms.steer_state.hold_course_cd != -1 && ground_steering)) {
        steering_output = fms_calc_nav_yaw_course();
    } else if (ground_steering) {
        steering_output = fms_calc_nav_yaw_ground();
    }

    /*
      now calculate rudder for the rudder
     */
    const float rudder_output = (float)fms_calc_nav_yaw_coordinated();

    if (!ground_steering) {
        // Not doing ground steering, output rudder on steering channel
        srv_channels_set_output_scaled(k_rudder, rudder_output);
        srv_channels_set_output_scaled(k_steering, rudder_output);

    } else if (!srv_channels_function_assigned(k_steering)) {
        // Ground steering active but no steering output configured, output steering on rudder channel
        srv_channels_set_output_scaled(k_rudder, steering_output);
        srv_channels_set_output_scaled(k_steering, steering_output);

    } else {
        // Ground steering with both steering and rudder channels
        srv_channels_set_output_scaled(k_rudder, rudder_output);
        srv_channels_set_output_scaled(k_steering, steering_output);
    }
}

/*
  calculate yaw control for ground steering with specific course
 */
int16_t fms_calc_nav_yaw_course(void)
{
    // holding a specific navigation course on the ground. Used in
    // auto-takeoff and landing
    int32_t bearing_error_cd = navctrl_bearing_error_cd(fms.nav_controller);
    int16_t steering = sctrl_get_steering_out_angle_error(&fms.steer_controller, bearing_error_cd);
    if (fms_stick_mixing_enabled()) {
        steering = rc_stick_mixing(fms.channel_rudder , steering);
    }
    return math_constrain_int16(steering, -4500, 4500);
}

/*
  calculate yaw control for ground steering
 */
int16_t fms_calc_nav_yaw_ground(void)
{
    if (fms.ahrs->gps_ground_speed < 1 && 
        math_flt_zero(fms_get_throttle_input(false)) &&
        fms.flight_stage != FWFS_TAKEOFF &&
        fms.flight_stage != FWFS_ABORT_LANDING) {
        // manual rudder control while still
        fms.steer_state.locked_course = false;
        fms.steer_state.locked_course_err = 0;
        return fms_rudder_input();
    }

    // if we haven't been steering for 1s then clear locked course
    const uint32_t now_ms = time_millis();
    if (now_ms - fms.steer_state.last_steer_ms > 1000) {
        fms.steer_state.locked_course = false;
    }
    fms.steer_state.last_steer_ms = now_ms;

    float steer_rate = (fms_rudder_input()/4500.0f) * fms.g.ground_steer_dps;
    if (fms.flight_stage == FWFS_TAKEOFF ||
        fms.flight_stage == FWFS_ABORT_LANDING) {
        steer_rate = 0;
    }
    if (!math_flt_zero(steer_rate)) {
        // pilot is giving rudder input
        fms.steer_state.locked_course = false;
    } else if (!fms.steer_state.locked_course) {
        // pilot has released the rudder stick or we are still - lock the course
        fms.steer_state.locked_course = true;
        if (fms.flight_stage != FWFS_TAKEOFF &&
            fms.flight_stage != FWFS_ABORT_LANDING) {
            fms.steer_state.locked_course_err = 0;
        }
    }

    int16_t steering;
    if (!fms.steer_state.locked_course) {
        // use a rate controller at the pilot specified rate
        steering = sctrl_get_steering_out_rate(&fms.steer_controller , steer_rate);
    } else {
        // use a error controller on the summed error
        int32_t yaw_error_cd = -ToDeg(fms.steer_state.locked_course_err)*100;
        steering = sctrl_get_steering_out_angle_error(&fms.steer_controller , yaw_error_cd);
    }

    return math_constrain_int16(steering, -4500, 4500);
}

/*
  calculate yaw control for coordinated flight
 */
int16_t fms_calc_nav_yaw_coordinated()
{
    const float speed_scaler = fms.surface_speed_scaler;
    bool disable_integrator = false;
    int16_t rudder_in = fms_rudder_input();

    int16_t commanded_rudder;
    bool using_rate_controller = false;

    // Received an external msg that guides yaw in the last 3 seconds?
    if (mode_is_guided_mode(fms.control_mode) &&
            fms.guided_state.last_forced_rpy_ms[2] > 0 &&
            time_millis() - fms.guided_state.last_forced_rpy_ms[2] < 3000) {
        commanded_rudder = fms.guided_state.forced_rpy_cd[2];
    } else if (fms.autotuning && fms.g.acro_yaw_rate > 0 && yctrl_rate_control_enabled(&fms.yaw_controller)) {
        // user is doing an AUTOTUNE with yaw rate control
        const float rudd_expo = fms_rudder_in_expo(true);
        const float yaw_rate = (rudd_expo/SERVO_MAX) * fms.g.acro_yaw_rate;
        // add in the corrdinated turn yaw rate to make it easier to fly while tuning the yaw rate controller
        const float coordination_yaw_rate = degrees(GRAVITY_MSS * tanf(radians(fms.nav_roll_cd*0.01f))/MAX(fms.aparm.airspeed_min,fms.smoothed_airspeed));
        commanded_rudder = yctrl_get_rate_out(&fms.yaw_controller , yaw_rate+coordination_yaw_rate,  speed_scaler, false);
        using_rate_controller = true;
    } else {
        if (fms.control_mode == (mode_base_t)&fms.mode_stabilize && rudder_in != 0) {
            disable_integrator = true;
        }

        commanded_rudder = yctrl_get_servo_out(&fms.yaw_controller, speed_scaler, disable_integrator);

        // add in rudder mixing from roll
        commanded_rudder += srv_channels_get_output_scaled(k_aileron) * fms.g.kff_rudder_mix;
        commanded_rudder += rudder_in;
    }

    if (!using_rate_controller) {
        /*
          When not running the yaw rate controller, we need to reset the rate
        */
        yctrl_reset_rate_pid(&fms.yaw_controller);
    }

    return math_constrain_int16(commanded_rudder, -4500, 4500);
}

/*------------------------------------test------------------------------------*/


